//
//  SPCSVExporter.m
//  sequel-pro
//
//  Created by Stuart Connolly (stuconnolly.com) on August 29, 2009.
//  Copyright (c) 2009 Stuart Connolly. All rights reserved.
//
//  Permission is hereby granted, free of charge, to any person
//  obtaining a copy of this software and associated documentation
//  files (the "Software"), to deal in the Software without
//  restriction, including without limitation the rights to use,
//  copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the
//  Software is furnished to do so, subject to the following
//  conditions:
//
//  The above copyright notice and this permission notice shall be
//  included in all copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
//  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
//  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
//  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
//  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
//  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
//  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
//  OTHER DEALINGS IN THE SOFTWARE.
//
//  More info at <https://github.com/sequelpro/sequelpro>

#import "SPCSVExporter.h"
#import "SPFileHandle.h"
#import "SPTableData.h"
#import "SPExportUtilities.h"
#import "SPExportFile.h"
#import "SPExportController.h"
#import "SPFunctions.h"

#import <SPMySQL/SPMySQL.h>

@implementation SPCSVExporter

@synthesize delegate;
@synthesize csvDataArray;
@synthesize csvTableName;
@synthesize csvOutputFieldNames;
@synthesize csvFieldSeparatorString;
@synthesize csvEnclosingCharacterString;
@synthesize csvEscapeString;
@synthesize csvLineEndingString;
@synthesize csvNULLString;
@synthesize csvTableData;

/**
 * Initialise an instance of SPCSVExporter using the supplied delegate.
 *
 * @param exportDelegate The exporter delegate
 *
 * @return The initialised instance
 */
- (instancetype)initWithDelegate:(NSObject<SPCSVExporterProtocol> *)exportDelegate
{
    if ((self = [super init])) {
        SPExportDelegateConformsToProtocol(exportDelegate, @protocol(SPCSVExporterProtocol));
        
        [self setDelegate:exportDelegate];
    }
    
    return self;
}

/**
 * Start the CSV export process.
 */
- (void)exportOperation
{
    NSMutableString *csvString     = [NSMutableString string];
    NSMutableString *csvCellString = [NSMutableString string];
    
    NSMutableArray *tableColumnNumericStatus = [NSMutableArray array];
    
    NSArray *csvRow = nil;
    NSScanner *csvNumericTester = nil;
    SPMySQLFastStreamingResult *streamingResult = nil;
    NSString *escapedEscapeString, *escapedFieldSeparatorString, *escapedEnclosingString, *escapedLineEndString, *dataConversionString;
    
    id csvCell;
    BOOL csvCellIsNumeric;
    BOOL quoteFieldSeparators = [[self csvEnclosingCharacterString] isEqualToString:@""];
    
    double lastProgressValue;
    NSUInteger i, totalRows, csvCellCount = 0;
    
    // Check to see if we have at least a table name or data array
    if ((![self csvTableName] && ![self csvDataArray]) ||
        ([[self csvTableName] isEqualToString:@""] && [[self csvDataArray] count] == 0))
    {
        return;
    }
    
    // Check that we have all the required info before starting the export
    if ((![self csvFieldSeparatorString]) ||
        (![self csvEscapeString]) ||
        (![self csvLineEndingString]))
    {
        return;
    }
    
    // Check that the CSV output options are not just empty strings
    if (([[self csvFieldSeparatorString] isEqualToString:@""]) ||
        ([[self csvEscapeString] isEqualToString:@""]) ||
        ([[self csvLineEndingString] isEqualToString:@""])) 
    {
        return;
    }
    
    // Inform the delegate that the export process is about to begin
    [delegate performSelectorOnMainThread:@selector(csvExportProcessWillBegin:) withObject:self waitUntilDone:NO];
    
    // Mark the process as running
    [self setExportProcessIsRunning:YES];
    
    lastProgressValue = 0;
    
    // Before the streaming query is started, build an array of numeric columns if a table
    // is being exported
    if ([self csvTableName] && (![self csvDataArray])) {
        NSDictionary *tableDetails = nil;
        
        // Determine whether the supplied table is actually a table or a view via the CREATE TABLE command, and get the table details
        SPMySQLResult *queryResult = [connection queryString:[NSString stringWithFormat:@"SHOW CREATE TABLE %@", [[self csvTableName] backtickQuotedString]]];
        [queryResult setReturnDataAsStrings:YES];
        
        if ([queryResult numberOfRows]) {
            id object = [[queryResult getRowAsDictionary] objectForKey:@"Create View"];
            
            tableDetails = [[NSDictionary alloc] initWithDictionary:(object) ? [[self csvTableData] informationForView:[self csvTableName]] : [[self csvTableData] informationForTable:[self csvTableName] fromDatabase:nil]];
        }
        
        // Retrieve the table details via the data class, and use it to build an array containing column numeric status
        for (NSDictionary *column in [tableDetails objectForKey:@"columns"])
        {
            NSString *tableColumnTypeGrouping = [column objectForKey:@"typegrouping"];
            
            [tableColumnNumericStatus addObject:[NSNumber numberWithBool:
                                                 ([tableColumnTypeGrouping isEqualToString:@"bit"]
                                                  || [tableColumnTypeGrouping isEqualToString:@"integer"]
                                                  || [tableColumnTypeGrouping isEqualToString:@"float"])
                                                 ]]; 
        }
    }
    
    // Make a streaming request for the data if the data array isn't set
    if ((![self csvDataArray]) && [self csvTableName]) {
        totalRows		= [[connection getFirstFieldFromQuery:[NSString stringWithFormat:@"SELECT COUNT(1) FROM %@", [[self csvTableName] backtickQuotedString]]] integerValue];
        streamingResult = [connection streamingQueryString:[NSString stringWithFormat:@"SELECT * FROM %@", [[self csvTableName] backtickQuotedString]] useLowMemoryBlockingStreaming:[self exportUsingLowMemoryBlockingStreaming]];
    }
    
    // Detect and restore special characters being used as terminating or line end strings
    NSMutableString *tempSeparatorString = [NSMutableString stringWithString:[self csvFieldSeparatorString]];
    
    // Escape tabs, line endings and carriage returns
    [tempSeparatorString replaceOccurrencesOfString:@"\\t" withString:@"\t"
                                            options:NSLiteralSearch
                                              range:NSMakeRange(0, [tempSeparatorString length])];
    
    [tempSeparatorString replaceOccurrencesOfString:@"\\n" withString:@"\n"
                                            options:NSLiteralSearch
                                              range:NSMakeRange(0, [tempSeparatorString length])];
    
    [tempSeparatorString replaceOccurrencesOfString:@"\\r" withString:@"\r"
                                            options:NSLiteralSearch
                                              range:NSMakeRange(0, [tempSeparatorString length])];
    
    // Set the new field separator string
    [self setCsvFieldSeparatorString:[NSString stringWithString:tempSeparatorString]];
    
    NSMutableString *tempLineEndString = [NSMutableString stringWithString:[self csvLineEndingString]];
    
    // Escape tabs, line endings and carriage returns
    [tempLineEndString replaceOccurrencesOfString:@"\\t" withString:@"\t"
                                          options:NSLiteralSearch
                                            range:NSMakeRange(0, [tempLineEndString length])];
    
    
    [tempLineEndString replaceOccurrencesOfString:@"\\n" withString:@"\n"
                                          options:NSLiteralSearch
                                            range:NSMakeRange(0, [tempLineEndString length])];
    
    [tempLineEndString replaceOccurrencesOfString:@"\\r" withString:@"\r"
                                          options:NSLiteralSearch
                                            range:NSMakeRange(0, [tempLineEndString length])];
    
    // Set the new line ending string
    [self setCsvLineEndingString:[NSString stringWithString:tempLineEndString]];
    
    // Set up escaped versions of strings for substitution within the loop
    escapedEscapeString         = [[self csvEscapeString] stringByAppendingString:[self csvEscapeString]];
    escapedFieldSeparatorString = [[self csvEscapeString] stringByAppendingString:[self csvFieldSeparatorString]];
    escapedEnclosingString      = [[self csvEscapeString] stringByAppendingString:[self csvEnclosingCharacterString]];
    escapedLineEndString        = [[self csvEscapeString] stringByAppendingString:[self csvLineEndingString]];
    
    // Set up the starting row; for supplied arrays, which include the column
    // headers as the first row, decide whether to skip the first row.
    NSUInteger currentRowIndex = 0;
    
    [csvString setString:@""];
    
    if ([self csvDataArray]) totalRows = [[self csvDataArray] count];
    if (([self csvDataArray]) && (![self csvOutputFieldNames])) currentRowIndex++;
    
    // Inform the delegate that we are about to start writing the data to disk
    [delegate performSelectorOnMainThread:@selector(csvExportProcessWillBeginWritingData:) withObject:self waitUntilDone:NO];
    
    while (1) 
    {
        if(self.exportOutputFile.fileHandleError != nil){
            SPMainQSync(^{
                [(SPExportController*)self->delegate cancelExportForFile:self->exportOutputFile.exportFilePath];
            });
            return;
        }
        // Check for cancellation flag
        if ([self isCancelled] ) {
            if (streamingResult) {
                [connection cancelCurrentQuery];
                [streamingResult cancelResultLoad];
            }
            
            return;
        }
        
        // Retrieve the next row from the supplied data, either directly from the array...
        BOOL forceNonNumericRow = NO;
        if ([self csvDataArray]) {
            csvRow = [[self csvDataArray] safeObjectAtIndex:currentRowIndex];
        }
        // Or by reading an appropriate row from the streaming result
        else {
            // If still requested to read the field names, get the field names
            if ([self csvOutputFieldNames]) {
                csvRow = [streamingResult fieldNames];
                [self setCsvOutputFieldNames:NO];
                forceNonNumericRow = YES;
            } 
            else {
                csvRow = [streamingResult getRowAsArray];
                
                if (!csvRow) break;
            }
        }
        
        // Get the cell count if we don't already have it stored
        if (!csvCellCount) csvCellCount = [csvRow count];
        
        [csvString setString:@""];
        
        for (i = 0 ; i < csvCellCount; i++) 
        {
            if(self.exportOutputFile.fileHandleError != nil){
                SPMainQSync(^{
                    [(SPExportController*)self->delegate cancelExportForFile:self->exportOutputFile.exportFilePath];
                });
                return;
            }
            // Check for cancellation flag
            if ([self isCancelled]) {
                return;
            }
            
            csvCell = [csvRow safeObjectAtIndex:i];
            
            // For NULL objects supplied from a queryResult, add an unenclosed null string as per prefs
            if ([csvCell isNSNull]) {
                [csvString appendString:[self csvNULLString]];
                
                if (i < (csvCellCount - 1)) [csvString appendString:[self csvFieldSeparatorString]];
                
                continue;
            }
            
            // Retrieve the contents of this cell
            if ([csvCell isKindOfClass:[NSData class]]) {
                dataConversionString = [[NSString alloc] initWithData:csvCell encoding:[self exportOutputEncoding]];
                
                if (dataConversionString == nil) {
                    dataConversionString = [[NSString alloc] initWithData:csvCell encoding:NSASCIIStringEncoding];
                }
                
                [csvCellString setString:[NSString stringWithString:dataConversionString]];
            }
            else if ([csvCell isKindOfClass:[SPMySQLGeometryData class]]) {
                [csvCellString setString:[csvCell wktString]];
            }
            else {
                [csvCellString setString:[csvCell description]];
            }
            
            // Add empty strings as a pair of enclosing characters.
            if ([csvCellString length] == 0) {
                [csvString appendString:[self csvEnclosingCharacterString]];
                [csvString appendString:[self csvEnclosingCharacterString]];
            }
            else {
                // is this the header row?
                if (forceNonNumericRow) {
                    csvCellIsNumeric = NO;
                }
                // If an array of bools supplying information as to whether the column is numeric has been supplied, use it.
                else if ([tableColumnNumericStatus count] > 0) {
                    csvCellIsNumeric = [[tableColumnNumericStatus safeObjectAtIndex:i] boolValue];
                } 
                // Otherwise, first test whether this cell contains data
                else if ([[csvRow safeObjectAtIndex:i] isKindOfClass:[NSData class]]) {
                    csvCellIsNumeric = NO;
                } 
                // Or fall back to testing numeric content via an NSScanner.
                else {
                    csvNumericTester = [NSScanner scannerWithString:csvCellString];
                    
                    csvCellIsNumeric = [csvNumericTester scanFloat:nil] && 
                    [csvNumericTester isAtEnd] && 
                    ([csvCellString characterAtIndex:0] != '0' || 
                     [csvCellString length] == 1 || 
                     ([csvCellString length] > 1 && 
                      [csvCellString characterAtIndex:1] == '.'));
                }
                
                // Escape any occurrences of the escaping character
                [csvCellString replaceOccurrencesOfString:[self csvEscapeString]
                                               withString:escapedEscapeString
                                                  options:NSLiteralSearch
                                                    range:NSMakeRange(0, [csvCellString length])];
                
                // Escape any occurrences of the enclosure string
                if (![[self csvEscapeString] isEqualToString:[self csvEnclosingCharacterString]]) {
                    [csvCellString replaceOccurrencesOfString:[self csvEnclosingCharacterString]
                                                   withString:escapedEnclosingString
                                                      options:NSLiteralSearch
                                                        range:NSMakeRange(0, [csvCellString length])];
                }
                
                // If the string isn't quoted or otherwise enclosed, escape occurrences of the field separators and line end character
                if (quoteFieldSeparators || csvCellIsNumeric) {
                    [csvCellString replaceOccurrencesOfString:[self csvFieldSeparatorString]
                                                   withString:escapedFieldSeparatorString
                                                      options:NSLiteralSearch
                                                        range:NSMakeRange(0, [csvCellString length])];
                    [csvCellString replaceOccurrencesOfString:[self csvLineEndingString]
                                                   withString:escapedLineEndString
                                                      options:NSLiteralSearch
                                                        range:NSMakeRange(0, [csvCellString length])];
                }
                
                // Write out the cell data by appending strings - this is significantly faster than stringWithFormat.
                if (csvCellIsNumeric) {
                    [csvString appendString:csvCellString];
                } 
                else {
                    [csvString appendString:[self csvEnclosingCharacterString]];
                    [csvString appendString:csvCellString];
                    [csvString appendString:[self csvEnclosingCharacterString]];
                }
            }
            
            if (i < ([csvRow count] - 1)) [csvString appendString:[self csvFieldSeparatorString]];
        }
        
        // Append the line ending to the string for this row, and record the length processed for pool flushing
        [csvString appendString:[self csvLineEndingString]];
        
        // Write it to the fileHandle
        [self writeString:csvString];
        
        currentRowIndex++;
        
        // Update the progress
        if (totalRows && (currentRowIndex * ([self exportMaxProgress] / totalRows)) > lastProgressValue) {
            
            double progress = (currentRowIndex * ([self exportMaxProgress] / totalRows));
            
            [self setExportProgressValue:progress];
            
            lastProgressValue = progress;
        }
        
        // Inform the delegate that the export's progress has been updated
        [delegate performSelectorOnMainThread:@selector(csvExportProcessProgressUpdated:) withObject:self waitUntilDone:NO];
        
        // If an array was supplied and we've processed all rows, break
        if ([self csvDataArray] && (totalRows == currentRowIndex)) break;
    }
    
    // Write data to disk
    [[[self exportOutputFile] exportFileHandle] synchronizeFile];
    
    // Mark the process as not running
    [self setExportProcessIsRunning:NO];
    
    // Inform the delegate that the export process is complete
    [delegate performSelectorOnMainThread:@selector(csvExportProcessComplete:) withObject:self waitUntilDone:NO];
}

@end
